package LibDL.eval.regression;

import LibDL.core.Dtype;
import LibDL.core.Scalar;
import LibDL.core.Tensor;
import LibDL.core.functional;

public class RSquaredEvaluator extends AbstractRegressionEvaluator{
    private boolean varianceWeighted = false;
    @Override
    protected double core(Tensor truth, Tensor pred) {
        TargetType type = checkTarget(truth, pred);
        if (type.equals(TargetType.MULTI_OUTPUT) && getOutputWeight() == null){
            setOutputWeight(functional.ones(truth.size(1)));
        }
        if (getSampleWeight() == null){
            setSampleWeight(functional.ones(truth.size(0)));
        }
        checkConsistentShape(truth,getSampleWeight(),getOutputWeight(),type);

        Tensor residualSS = weightedSum(truth.sub(pred).pow(new Scalar(2)),
                getSampleWeight(),true,0);
        Tensor totalSS = weightedSum(
                truth.sub(weightedSum(truth,getSampleWeight(),true,0).unsqueeze(0))
                        .pow(new Scalar(2))
                ,getSampleWeight(),true,0
        );
        if (residualSS.sizes().size() == 0){
            return 1 - residualSS.item().to_double()/totalSS.item().to_double();
        }
        Tensor nonzeroResidual = functional.equal_indices(residualSS,new Scalar(0))
                .to(Dtype.FLOAT64).sub(new Scalar(1)).mul(-1);
        Tensor zeroTotal = functional.equal_indices(totalSS,new Scalar(0)).to(Dtype.FLOAT64);
        Tensor valid = functional.equal_indices(totalSS.mul(residualSS),new Scalar(0))
                .to(Dtype.FLOAT64).sub(new Scalar(1)).mul(-1);
        Tensor rawVal = functional.ones(residualSS.size(0));
        rawVal = rawVal.sub(
                residualSS.div(totalSS.add(zeroTotal.to(Dtype.FLOAT64))).mul(valid)
        );
        rawVal = rawVal.sub(
                functional.equal_indices(nonzeroResidual.add(zeroTotal),new Scalar(2))
                        .to(Dtype.FLOAT64)
        );

        if (isVarianceWeighted()) {
            return weightedSum(rawVal,totalSS,true,0).item().to_double();
        }
        return weightedSum(rawVal,getOutputWeight(),true,0).item().to_double();
    }

    public boolean isVarianceWeighted() {
        return varianceWeighted;
    }

    public RSquaredEvaluator setVarianceWeighted(boolean varianceWeighted) {
        this.varianceWeighted = varianceWeighted;
        return this;
    }
}
